"""
接口：express
描述：获取上市公司业绩快报
权限：用户需要至少2000积分才可以调取，具体请参阅积分获取办法
"""

import pandas as pd
import time
from app_config import get_pro, get_engine_ts

pro = get_pro()
engine_ts = get_engine_ts()


def read_data():
    sql = """SELECT * FROM express"""
    df = pd.read_sql_query(sql, engine_ts)
    return df


def write_data(df):
    res = df.to_sql('express', engine_ts, index=False, if_exists='append', chunksize=5000)
    print(res)


def get_data(_ts_code):
    df = pro.express(ts_code=_ts_code,
                     fields='ts_code,ann_date,end_date,revenue,operate_profit,total_profit,n_income,total_assets')
    return df


def get_express(_ts_code):
    for _ in range(10):
        try:
            df = get_data(_ts_code)
            # 创建新的列 'id'，将 'ts_code' 和 'trade_date' 两列拼接起来
            df['id'] = df['ts_code'].astype(str) + "_" + df['end_date'].astype(str)
            write_data(df)
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            time.sleep(10)
        else:
            return
    print("error error error error error" + _ts_code)


if __name__ == '__main__':

    from _stock_basic import get_data as get_stock_basic

    # 获取20230601～20240531之间所有有交易的日期
    sql = """SELECT * FROM stock_basic"""
    basic = pd.read_sql_query(sql, engine_ts)

    for index, row in basic.iterrows():  #
        # 每天的交易日期
        ts_code = row["ts_code"]
        print("股票代码: " + str(ts_code) + str(row["name"]))
        get_express(ts_code)
